# !/usr/bin/python
# -*- coding: utf-8 -*-

import os
import shutil
import warnings

import fitz
import win32api
import win32con
from cnocr import CnOcr
from cnocr.utils import read_img


class Pdf2Txt:

    def __init__(self):
        pass

    # ---把pdf每页都转换成一张图片--- #
    def pdf_to_img(self, input_pdf_name, output_file):
        pdf = fitz.open(input_pdf_name + '.pdf')
        page = pdf.page_count  # 读取pdf的页数
        for i in range(0, page):
            page_fn = pdf.load_page(i)  # 读取第i页
            pix = page_fn.get_pixmap()
            pix.save(os.path.join(output_file, str(i) + '.png'))  # 从0开始按序号另存页数为图片 v5.0修改：用函数链接输出路径，避免不同系统的分隔符差异带来的报错

    # ---识别图片中的文字并且return--- #
    def img_ocr(self, input_pic):
        ocr = CnOcr()
        img_fp = input_pic  # 输入图片路径
        img = read_img(img_fp)  # 读取图片
        res = ocr.ocr(img)  # res是list（列表）
        len_res = len(res)  # 文本的行数
        text = ""
        for i in range(0, len_res):
            line_tuple = res[i]  # line_tuple是tuple（元组）
            len_line_tuple = len(line_tuple[0])
            s = ""
            single_line = s.join(line_tuple[0])
            text += single_line
        return text  # text为识别出来的文本，此处return可以给下面写入txt

    # ---将识别出来的文字写入txt--- #
    def write_to_txt(self, input_pdf_name, text):
        with open(input_pdf_name + '.txt', 'a', encoding='utf-8') as file:  # 用a模式是续写而不是w(重新写入),编码为utf-8
            file.write(text)
            file.write('\n')  # 写完一张图片里面的内容就换行

    # ---获取当前路径下所有pdf文件--- #
    def get_pdf_list(self):
        pdf_list = [x for x in os.listdir('.') if os.path.isfile(x) and os.path.splitext(x)[1] == '.pdf']
        return pdf_list

    # ---忽略报错--- #
    def ignore_warning(self):
        warnings.filterwarnings("ignore")

    # ---连接路径--- #
    def join_path(self, path1, path2):
        join_path = os.path.join(path1, path2)
        return join_path

    # ---创建目录--- #
    def mk_dir(self, filename):
        os.mkdir(filename)

    # ---删除目录--- #
    def rm_dir(self, filename):
        shutil.rmtree(filename)

    # ---弹窗提示完成--- #
    def finish(self):
        # ctypes.windll.user32.MessageBoxW(0, u'已完成PDF转化', u'已完成', 0)
        win32api.MessageBox(0, "已完成PDF转化", "已完成", win32con.MB_ICONASTERISK)
        # os.system("msg * "ok"")

    def start_ocr(self):
        self.ignore_warning()
        pdf_list = self.get_pdf_list()
        for i in range(0, len(pdf_list)):
            name = os.path.splitext(pdf_list[i])[0]
            input_pdf_name = self.join_path(os.path.curdir, name)
            output_file = self.join_path(os.path.curdir, 'pdf_2_image')
            input_file = output_file
            self.mk_dir(output_file)
            self.pdf_to_img(input_pdf_name, output_file)
            pic = os.listdir(input_file)
            for pic_fn in pic:
                text = self.img_ocr(self.join_path(input_file, pic_fn))
                self.write_to_txt(input_pdf_name, text)
            self.rm_dir(output_file)
        self.finish()
